import cv2
import argparse
import os
import numpy as np
import sys
from pathlib import Path
 
sys.path.append(str(Path(__file__).resolve().parent.parent))

from file_util.file_util import GetAllSubFile
import xml.etree.ElementTree as ET

def draw_image_with_anno(new_img_file,new_anno_file):
		try:
			tree = ET.parse(new_anno_file)
			root = tree.getroot() 
		except:
			print(new_anno_file + " is not exist.")
			return
		img = cv2.imread(new_img_file)
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)
				cv2.rectangle(img,(xmin,ymin),(xmax,ymax),(0,0,255),5)
		cv2.imwrite(new_img_file,img)


def update_anno(anno_file,diff_x, diff_y,w,h, new_anno_file):
		try:
			tree = ET.parse(anno_file)
			root = tree.getroot() 
		except:
			print(anno_file + " is not exist.")
			return
		remove_objects = []
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)
				print(object[0].text,"xmin=",xmin,"ymin=",ymin,"xmax=",xmax,"ymax=",ymax,"diff_x=",diff_x,"diff_y=",diff_y,"w=",w,"h=",h)
				new_xmin = xmin + diff_x
				new_ymin = ymin + diff_y
				new_xmax = xmax + diff_x
				new_ymax = ymax + diff_y
				if (new_xmax < 0) or (new_ymax < 0):
					print(object[0].text)
					#root.remove(object)
					remove_objects.append(object)
					continue
				if (new_xmin > w) or (new_ymin > h):
					print(object[0].text)
					#root.remove(object)
					remove_objects.append(object)
					continue
				if new_xmin < 0:
					new_xmin = 0
				if new_ymin < 0 :
					new_ymin = 0
				if new_xmax > w :
					new_xmax = w
				if new_ymax > h :
					new_ymax = h
				print(object[0].text,"new_xmin=",new_xmin,"new_ymin=",new_ymin,"new_xmax=",new_xmax,"new_ymax=",new_ymax)
				object[4][0].text = str(new_xmin)
				object[4][1].text = str(new_ymin)
				object[4][2].text = str(new_xmax)
				object[4][3].text = str(new_ymax)
		for remove_object in remove_objects:
			root.remove(remove_object)
		tree.write(new_anno_file,encoding="utf-8", xml_declaration=True)
		return 

def img_shift(opt,shift_list):
	img_files = GetAllSubFile(opt.root_data)
	for img_file in img_files:
		img = cv2.imread(img_file)
		h = img.shape[0]
		w = img.shape[1]
		path, img_file_name = os.path.split(img_file)
		img_name = img_file_name.rsplit('.',1)[0]
		#draw_totation(img_file,os.path.join(opt.anon_dir,img_name + '.xml'))
		anno_file = os.path.join(opt.anon_dir,img_name + '.xml')
		for t_x, t_y in shift_list:
			M = np.float32([[1, 0, t_x], [0, 1, t_y]])
			shifted_img = cv2.warpAffine(img, M, (w, h),borderValue=(114,114,114))
			new_img_name = img_name+"_shift_" + str(t_x) + "_" + str(t_y)
			new_anno_file = os.path.join(opt.anon_dir,new_img_name + '.xml')
			new_img_file = os.path.join(path,new_img_name +".jpg")
			cv2.imwrite(new_img_file,shifted_img)
			print(new_anno_file)
			update_anno(anno_file,int(t_x), int(t_y), w, h, new_anno_file)
			#draw_image_with_anno(new_img_file,new_anno_file)

if __name__ == '__main__':
   parser = argparse.ArgumentParser()
   parser.add_argument('--root_data', type=str, default='./data', help='root data dir.')
   parser.add_argument('--anon_dir', type=str, default='./anno', help='anno dir.')
   opt = parser.parse_args()
   opt.anon_dir = opt.anon_dir.strip()
   #if x and y great than 0, move toward bottom right
   shift_list = [(50,50),(200,200),(50,-50),(200,-200)]
   shift_list.extend([(-50,50),(-200,200)])
   shift_list.extend([(-50,-50),(-200,-200)])
   #shift_list = [(200, 200)]
   #8
   img_shift(opt,shift_list)

